"""
向量模型工厂类 - 策略模式实现
"""

import os
from functools import lru_cache
from typing import Dict, Type, Optional, Any

from dotenv import load_dotenv

from config.logging_config import get_logger
from .BaseEmb import BaseEmb, ModelNotFoundError, ConfigurationError
from .DashScopeEmb import DashScopeEmb
from .ZhiPuEmb import ZhiPuEmb

# ChromaDB 相关导入
try:
    from chromadb import EmbeddingFunction, Documents, Embeddings

    CHROMADB_AVAILABLE = True
except ImportError:
    CHROMADB_AVAILABLE = False
    EmbeddingFunction = object  # 占位符

logger = get_logger(__name__)

# 加载环境变量
load_dotenv()


class ChromaEmbeddingFunction(EmbeddingFunction):
    """ChromaDB嵌入函数适配器 - 通用实现"""

    def __init__(self, emb_model: BaseEmb):
        """
        初始化嵌入函数
        
        Args:
            emb_model: 向量模型实例
        """
        self.emb_model = emb_model
        self.model_name = emb_model.model_name
        logger.info(f"ChromaDB嵌入函数初始化成功，模型: {self.model_name}")

    def __call__(self, input: Documents) -> Embeddings:
        """
        执行嵌入计算
        
        Args:
            input: 输入文档列表
            
        Returns:
            嵌入向量列表
        """
        try:
            if isinstance(input, str):
                input = [input]

            # 使用向量模型进行向量化
            embeddings = self.emb_model.encode(input)
            return embeddings

        except Exception as e:
            logger.error(f"嵌入计算失败: {e}")
            raise


class EmbFactory:
    """向量模型工厂类，负责根据emb_code路由到对应的实现类"""

    # 注册的向量模型实现类
    _registry: Dict[str, Type[BaseEmb]] = {}

    # 模型配置映射
    _model_configs: Dict[str, Dict[str, Any]] = {
        # 智谱AI模型配置
        "zhipu-embedding-2": {
            "class": "ZhiPuEmb",
            "model_name": "embedding-2",
            "provider": "zhipu"
        },
        "zhipu-embedding-3": {
            "class": "ZhiPuEmb",
            "model_name": "embedding-3",
            "provider": "zhipu"
        },
        "text-embedding-3-large": {
            "class": "ZhiPuEmb",
            "model_name": "text-embedding-3-large",
            "provider": "zhipu"
        },
        "text-embedding-3-small": {
            "class": "ZhiPuEmb",
            "model_name": "text-embedding-3-small",
            "provider": "zhipu"
        },
        "dashscope-text-embedding-v3": {
            "class": "DashScopeEmb",
            "model_name": "text-embedding-v3",
            "provider": "dashscope"
        },
        "dashscope-text-embedding-v4": {
            "class": "DashScopeEmb",
            "model_name": "text-embedding-v4",
            "provider": "dashscope"
        },
        # 可以继续添加其他提供商的模型
        # "openai-text-embedding-3-large": {
        #     "class": "OpenAIEmb",
        #     "model_name": "text-embedding-3-large",
        #     "provider": "openai"
        # },
    }

    @classmethod
    def register_model(cls, emb_code: str, model_class: Type[BaseEmb]) -> None:
        """
        注册向量模型实现类
        
        Args:
            emb_code: 模型代码
            model_class: 模型实现类
        """
        cls._registry[emb_code] = model_class
        logger.info(f"注册向量模型: {emb_code}")

    @classmethod
    def _auto_register_models(cls) -> None:
        """自动注册内置模型"""
        if not cls._registry:
            # 注册智谱AI模型
            cls.register_model("ZhiPuEmb", ZhiPuEmb)

            # 注册阿里百炼DashScope模型
            cls.register_model("DashScopeEmb", DashScopeEmb)

            # 可以在这里注册其他模型
            # cls.register_model("OpenAIEmb", OpenAIEmb)
            # cls.register_model("HuggingFaceEmb", HuggingFaceEmb)

    @classmethod
    def get_supported_models(cls) -> Dict[str, Dict[str, Any]]:
        """
        获取支持的模型列表
        
        Returns:
            支持的模型配置字典
        """
        return cls._model_configs.copy()

    @classmethod
    def is_supported(cls, emb_code: str) -> bool:
        """
        检查是否支持指定的模型代码
        
        Args:
            emb_code: 模型代码
            
        Returns:
            是否支持
        """
        return emb_code in cls._model_configs

    @classmethod
    def get_model_config(cls, emb_code: str) -> Optional[Dict[str, Any]]:
        """
        获取模型配置
        
        Args:
            emb_code: 模型代码
            
        Returns:
            模型配置字典，如果不存在返回None
        """
        return cls._model_configs.get(emb_code)

    @classmethod
    @lru_cache(maxsize=32)
    def create_embedding_model(cls, emb_code: str, **kwargs) -> BaseEmb:
        """
        创建向量模型实例
        
        Args:
            emb_code: 模型代码
            **kwargs: 模型配置参数
            
        Returns:
            向量模型实例
            
        Raises:
            ModelNotFoundError: 模型不存在
            ConfigurationError: 配置错误
        """
        # 自动注册模型
        cls._auto_register_models()

        # 检查模型是否支持
        if not cls.is_supported(emb_code):
            supported_models = list(cls._model_configs.keys())
            raise ModelNotFoundError(
                f"不支持的向量模型: {emb_code}。"
                f"支持的模型: {', '.join(supported_models)}"
            )

        # 获取模型配置
        model_config = cls.get_model_config(emb_code)
        if not model_config:
            raise ModelNotFoundError(f"模型配置不存在: {emb_code}")

        # 获取模型实现类
        class_name = model_config["class"]
        if class_name not in cls._registry:
            raise ModelNotFoundError(f"模型实现类未注册: {class_name}")

        model_class = cls._registry[class_name]

        # 准备模型参数
        model_params = {
            "model_name": model_config["model_name"],
            **kwargs
        }

        # 根据提供商添加特定配置
        provider = model_config.get("provider")
        if provider == "zhipu":
            # 智谱AI需要API密钥
            api_key = kwargs.get("api_key") or os.getenv("ZHIPU_API_KEY")
            if not api_key:
                raise ConfigurationError(
                    "智谱AI模型需要API密钥。请在参数中提供api_key或设置环境变量ZHIPU_API_KEY"
                )
            model_params["api_key"] = api_key
        elif provider == "dashscope":
            # 阿里百炼DashScope需要API密钥
            api_key = kwargs.get("api_key") or os.getenv("DASHSCOPE_API_KEY")
            if not api_key:
                raise ConfigurationError(
                    "阿里百炼DashScope模型需要API密钥。请在参数中提供api_key或设置环境变量DASHSCOPE_API_KEY"
                )
            model_params["api_key"] = api_key
            # 传递dimensions参数（如果有）
            if "dimensions" in kwargs:
                model_params["dimensions"] = kwargs["dimensions"]

        # 创建模型实例
        try:
            model_instance = model_class(**model_params)
            logger.info(f"成功创建向量模型: {emb_code} ({class_name})")
            return model_instance
        except Exception as e:
            logger.error(f"创建向量模型失败: {emb_code} - {str(e)}")
            raise ConfigurationError(f"创建向量模型失败: {str(e)}")

    @classmethod
    def create_from_config(cls, config: Dict[str, Any]) -> BaseEmb:
        """
        从配置字典创建向量模型
        
        Args:
            config: 配置字典，必须包含emb_code字段
            
        Returns:
            向量模型实例
        """
        emb_code = config.get("emb_code")
        if not emb_code:
            raise ConfigurationError("配置中缺少emb_code字段")

        # 移除emb_code，其余参数传递给模型
        model_params = {k: v for k, v in config.items() if k != "emb_code"}

        return cls.create_embedding_model(emb_code, **model_params)

    @classmethod
    def test_model(cls, emb_code: str, **kwargs) -> bool:
        """
        测试模型是否可用
        
        Args:
            emb_code: 模型代码
            **kwargs: 模型配置参数
            
        Returns:
            模型是否可用
        """
        try:
            model = cls.create_embedding_model(emb_code, **kwargs)
            if hasattr(model, 'test_connection'):
                return model.test_connection()
            else:
                # 如果没有test_connection方法，尝试简单的向量化测试
                test_vector = model.get_emb("测试")
                return len(test_vector) > 0
        except Exception as e:
            logger.error(f"模型测试失败: {emb_code} - {str(e)}")
            return False

    @classmethod
    def clear_cache(cls) -> None:
        """清除缓存"""
        cls.create_embedding_model.cache_clear()
        logger.info("向量模型缓存已清除")

    @classmethod
    def create_embedding_function(cls, emb_code: str, **kwargs) -> Optional[EmbeddingFunction]:
        """
        创建ChromaDB嵌入函数
        
        Args:
            emb_code: 模型代码
            **kwargs: 模型配置参数
            
        Returns:
            ChromaDB嵌入函数实例，如果ChromaDB不可用则返回None
            
        Raises:
            ConfigurationError: 配置错误
            ModelNotFoundError: 模型不存在
        """
        if not CHROMADB_AVAILABLE:
            logger.warning("ChromaDB不可用，无法创建嵌入函数")
            return None

        try:
            # 创建向量模型实例
            emb_model = cls.create_embedding_model(emb_code, **kwargs)

            # 创建ChromaDB嵌入函数适配器
            embedding_function = ChromaEmbeddingFunction(emb_model)

            logger.info(f"成功创建ChromaDB嵌入函数: {emb_code}")
            return embedding_function

        except Exception as e:
            logger.error(f"创建ChromaDB嵌入函数失败: {emb_code} - {str(e)}")
            raise

    @classmethod
    def create_embedding_function_from_config(cls, config: Dict[str, Any]) -> Optional[EmbeddingFunction]:
        """
        从配置字典创建ChromaDB嵌入函数
        
        Args:
            config: 配置字典，必须包含emb_code字段
            
        Returns:
            ChromaDB嵌入函数实例
        """
        emb_code = config.get("emb_code")
        if not emb_code:
            raise ConfigurationError("配置中缺少emb_code字段")

        # 移除emb_code，其余参数传递给模型
        model_params = {k: v for k, v in config.items() if k != "emb_code"}

        return cls.create_embedding_function(emb_code, **model_params)

    @classmethod
    def get_default_embedding_function(cls,
                                       emb_code: str = "zhipu-embedding-3",
                                       fallback_to_default: bool = True) -> EmbeddingFunction:
        """
        获取默认的嵌入函数
        
        Args:
            emb_code: 首选的模型代码
            fallback_to_default: 是否在失败时回退到默认函数
            
        Returns:
            嵌入函数实例
        """
        if not CHROMADB_AVAILABLE:
            raise ConfigurationError("ChromaDB不可用")

        try:
            # 尝试创建指定的嵌入函数
            embedding_function = cls.create_embedding_function(emb_code)
            if embedding_function:
                return embedding_function
        except Exception as e:
            logger.warning(f"创建指定嵌入函数失败: {emb_code} - {str(e)}")

        # 回退到默认函数
        if fallback_to_default:
            try:
                from chromadb.utils import embedding_functions
                logger.info("回退到ChromaDB默认嵌入函数")
                return embedding_functions.DefaultEmbeddingFunction()
            except Exception as e:
                logger.error(f"创建默认嵌入函数失败: {e}")
                raise ConfigurationError("无法创建任何嵌入函数")
        else:
            raise ConfigurationError(f"无法创建嵌入函数: {emb_code}")


# 便捷函数
def create_embedding_model(emb_code: str, **kwargs) -> BaseEmb:
    """
    创建向量模型的便捷函数
    
    Args:
        emb_code: 模型代码
        **kwargs: 模型配置参数
        
    Returns:
        向量模型实例
    """
    return EmbFactory.create_embedding_model(emb_code, **kwargs)


def get_supported_models() -> Dict[str, Dict[str, Any]]:
    """
    获取支持的模型列表的便捷函数
    
    Returns:
        支持的模型配置字典
    """
    return EmbFactory.get_supported_models()


def test_embedding_model(emb_code: str, **kwargs) -> bool:
    """
    测试向量模型的便捷函数
    
    Args:
        emb_code: 模型代码
        **kwargs: 模型配置参数
        
    Returns:
        模型是否可用
    """
    return EmbFactory.test_model(emb_code, **kwargs)


def create_embedding_function(emb_code: str, **kwargs) -> Optional[EmbeddingFunction]:
    """
    创建ChromaDB嵌入函数的便捷函数
    
    Args:
        emb_code: 模型代码
        **kwargs: 模型配置参数
        
    Returns:
        ChromaDB嵌入函数实例
    """
    return EmbFactory.create_embedding_function(emb_code, **kwargs)


def get_default_embedding_function(emb_code: str = "zhipu-embedding-3",
                                   fallback_to_default: bool = True) -> EmbeddingFunction:
    """
    获取默认嵌入函数的便捷函数
    
    Args:
        emb_code: 首选的模型代码
        fallback_to_default: 是否在失败时回退到默认函数
        
    Returns:
        嵌入函数实例
    """
    return EmbFactory.get_default_embedding_function(emb_code, fallback_to_default)
